import gradio as gr
import os
import json
import requests
import time
from datetime import datetime
import pandas as pd
from image_generator import ImageGenerator
from video_generator import VideoGenerator

def load_config():
    config_file = "config.json"
    if os.path.exists(config_file):
        with open(config_file, 'r') as f:
            return json.load(f)
    return {"api_key": "", "api_url": ""}

def save_config(api_key, api_url):
    config = {
        "api_key": api_key,
        "api_url": api_url
    }
    with open("config.json", 'w') as f:
        json.dump(config, f)
    return "配置已保存"

def text_to_image(api_key, prompt):
    if not api_key:
        return None, "请先配置API Key"
    
    if not prompt:
        return None, "请输入提示词"
    
    try:
        generator = ImageGenerator(api_key)
        result = generator.generate_image(prompt)
        
        # 下载图片
        response = requests.get(result["url"])
        image_path = f"generated_images/generated_image_{int(time.time())}.jpg"
        os.makedirs("generated_images", exist_ok=True)
        
        with open(image_path, "wb") as f:
            f.write(response.content)
        
        return image_path, json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return None, f"请求异常: {str(e)}"

def image_to_image(api_key, prompt, image):
    if not api_key:
        return None, "请先配置API Key"
    
    if not image:
        return None, "请选择图片"
    
    if not prompt:
        return None, "请输入提示词"
    
    try:
        generator = ImageGenerator(api_key)
        
        # 保存上传的图片
        image_path = f"generated_images/uploaded_image_{int(time.time())}.jpg"
        image.save(image_path)
        
        # 上传图片并获取URL
        with open(image_path, "rb") as img_file:
            import base64
            img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
            image_url = f"data:image/jpeg;base64,{img_base64}"
        
        # 调用图像生成器
        result = generator.generate_image_from_image(prompt, image_url)
        
        # 下载结果图片
        response = requests.get(result["url"])
        result_image_path = f"generated_images/generated_image_{int(time.time())}.jpg"
        os.makedirs("generated_images", exist_ok=True)
        
        with open(result_image_path, "wb") as f:
            f.write(response.content)
        
        return result_image_path, json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return None, f"请求异常: {str(e)}"

def text_to_video(api_key, prompt):
    if not api_key:
        return None, "请先配置API Key", ""
    
    if not prompt:
        return None, "请输入提示词", ""
    
    try:
        generator = VideoGenerator(api_key=api_key)
        result = generator.generate_video(prompt)
        
        return None, json.dumps(result, indent=2, ensure_ascii=False), result["task_id"]
    except Exception as e:
        return None, f"请求异常: {str(e)}", ""

def check_video_status(api_key, task_id):
    if not api_key:
        return None, "请先配置API Key", ""
    
    if not task_id:
        return None, "请先生成视频任务", ""
    
    try:
        generator = VideoGenerator(api_key=api_key)
        result = generator.get_video_generation_status(task_id)
        
        # 如果视频已生成完成，下载并显示视频
        video_path = None
        if result["status"] == "succeeded":
            # 从结果中获取视频URL
            video_url = None
            
            # 检查不同的可能结构
            if "content" in result and result["content"]:
                content = result["content"]
                # 如果content是列表，检查列表中的元素
                if isinstance(content, list) and len(content) > 0:
                    first_item = content[0]
                    if isinstance(first_item, dict) and "video_url" in first_item:
                        video_url = first_item["video_url"]
                # 如果content是字典，直接检查video_url
                elif isinstance(content, dict) and "video_url" in content:
                    video_url = content["video_url"]
            
            # 如果在content中没找到，检查result
            if not video_url and "result" in result and result["result"]:
                res = result["result"]
                if isinstance(res, dict) and "video_url" in res:
                    video_url = res["video_url"]
            
            if video_url:
                # 下载视频
                os.makedirs("generated_videos", exist_ok=True)
                video_path = f"generated_videos/generated_video_{int(time.time())}.mp4"
                response = requests.get(video_url)
                response.raise_for_status()
                
                with open(video_path, "wb") as f:
                    f.write(response.content)
        
        return video_path, json.dumps(result, indent=2, ensure_ascii=False), task_id
    except Exception as e:
        return None, f"请求异常: {str(e)}", ""

def auto_check_video_status(api_key, task_id, progress=gr.Progress()):
    """自动检查视频生成状态，每10秒检查一次"""
    if not api_key:
        return None, "请先配置API Key", ""
    
    if not task_id:
        return None, "请先生成视频任务", ""
    
    # 显示加载弹窗
    progress(0, desc="视频生成中...")
    
    while True:
        try:
            generator = VideoGenerator(api_key=api_key)
            result = generator.get_video_generation_status(task_id)
            
            # 如果视频已生成完成，下载并显示视频
            if result["status"] == "succeeded":
                # 从结果中获取视频URL
                video_url = None
                
                # 检查不同的可能结构
                if "content" in result and result["content"]:
                    content = result["content"]
                    # 如果content是列表，检查列表中的元素
                    if isinstance(content, list) and len(content) > 0:
                        first_item = content[0]
                        if isinstance(first_item, dict) and "video_url" in first_item:
                            video_url = first_item["video_url"]
                    # 如果content是字典，直接检查video_url
                    elif isinstance(content, dict) and "video_url" in content:
                        video_url = content["video_url"]
                
                # 如果在content中没找到，检查result
                if not video_url and "result" in result and result["result"]:
                    res = result["result"]
                    if isinstance(res, dict) and "video_url" in res:
                        video_url = res["video_url"]
                
                if video_url:
                    # 下载视频
                    os.makedirs("generated_videos", exist_ok=True)
                    video_path = f"generated_videos/generated_video_{int(time.time())}.mp4"
                    response = requests.get(video_url)
                    response.raise_for_status()
                    
                    with open(video_path, "wb") as f:
                        f.write(response.content)
                    
                    return video_path, json.dumps(result, indent=2, ensure_ascii=False), task_id
                else:
                    # 继续等待
                    progress(0.8, desc="视频生成完成，正在获取下载链接...")
                    time.sleep(10)
                    continue
            elif result["status"] == "failed":
                return None, f"视频生成失败: {json.dumps(result, indent=2, ensure_ascii=False)}", task_id
            else:
                # 更新进度信息
                progress(0.5, desc="视频生成中，请稍候...")
                # 等待10秒后再次检查
                time.sleep(10)
        except Exception as e:
            return None, f"请求异常: {str(e)}", ""



def image_to_video_first_frame(api_key, prompt, image):
    print("Function called")
    print("===== DEBUG: image_to_video_first_frame 函数开始执行 =====")
    print(f"函数开始执行，api_key: {api_key}, prompt: {prompt}, image: {image}")
    print(f"当前时间戳: {int(time.time())}"
)
    if not api_key:
        print("API Key未配置")
        return None, "请先配置API Key"
    
    if not prompt:
        print("提示词未输入")
        return None, "请输入提示词"
    
    if image is None:
        print("图片未上传")
        return None, "请上传图片"
    
    try:
        # 保存上传的图片
        image_path = f"generated_images/uploaded_image_{int(time.time())}.jpg"
        os.makedirs("generated_images", exist_ok=True)
        image.save(image_path)
        print(f"图片已保存到: {image_path}")

        # 上传图片并获取URL
        with open(image_path, "rb") as img_file:
            import base64
            img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
            image_url = f"data:image/jpeg;base64,{img_base64}"

        # print(f"获取到url为{image_url}")
        # 调用视频生成API（首帧功能）
        video_generator = VideoGenerator(api_key)
        result = video_generator.generate_video_from_image(prompt, image_url)
        print(f"视频生成结果: {result}")
        print("===== DEBUG: image_to_video_first_frame 函数执行结束 =====")
        
        return None, json.dumps(result, indent=2, ensure_ascii=False), result["task_id"]
    except Exception as e:
        print(f"请求异常: {str(e)}")
        return None, f"请求异常: {str(e)}", ""

def load_badcase_records():
    badcase_file = "badcase_records.xlsx"
    if os.path.exists(badcase_file):
        df = pd.read_excel(badcase_file)
        return df
    return pd.DataFrame(columns=["时间", "提示词", "输入图片", "输出结果", "错误信息"])

def save_badcase_record(prompt, input_image, output_result, error_message):
    # 创建新的记录
    new_record = {
        "时间": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "提示词": prompt,
        "输入图片": input_image if input_image else "",
        "输出结果": output_result if output_result else "",
        "错误信息": error_message
    }
    
    # 读取现有的记录
    df = load_badcase_records()
    
    # 添加新记录
    new_df = pd.DataFrame([new_record])
    df = pd.concat([df, new_df], ignore_index=True)
    
    # 保存到Excel文件
    df.to_excel("badcase_records.xlsx", index=False)
    
    return "错误记录已保存"

def main():
    config = load_config()
    
    with gr.Blocks(title="马帮媒体生成POC测试工具") as demo:
        gr.Markdown("# 马帮媒体生成POC测试工具")
        
        with gr.Tab("API配置"):
            api_key = gr.Textbox(label="API Key", value=config["api_key"], type="password")
            api_url = gr.Textbox(label="API URL", value=config["api_url"])
            save_config_btn = gr.Button("保存配置")
            config_status = gr.Textbox(label="配置状态", interactive=False)
            
            save_config_btn.click(
                fn=save_config,
                inputs=[api_key, api_url],
                outputs=config_status
            )
        
        with gr.Tab("文字生成图片"):
            with gr.Row():
                with gr.Column():
                    text_prompt = gr.Textbox(label="提示词", lines=4)
                    text_to_image_btn = gr.Button("生成图片")
                    text_report_badcase_btn = gr.Button("记录错误")
                with gr.Column():
                    text_result_image = gr.Image(label="生成的图片")
                    text_result_text = gr.Textbox(label="结果详情", lines=10)
            
            text_to_image_btn.click(
                fn=text_to_image,
                inputs=[api_key, text_prompt],
                outputs=[text_result_image, text_result_text]
            )
            
            text_report_badcase_btn.click(
                fn=save_badcase_record,
                inputs=[text_prompt, gr.State(None), text_result_text, text_result_text]
            ).then(
                fn=lambda: gr.Info("错误记录已保存"),
                inputs=[],
                outputs=[]
            )
        
        with gr.Tab("图片生成图片"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(label="选择图片", type="pil")
                    image_prompt = gr.Textbox(label="提示词", lines=4)
                    image_to_image_btn = gr.Button("生成图片")
                    image_report_badcase_btn = gr.Button("记录错误")
                with gr.Column():
                    image_result_image = gr.Image(label="生成的图片")
                    image_result_text = gr.Textbox(label="结果详情", lines=10)
            
            image_to_image_btn.click(
                fn=image_to_image,
                inputs=[api_key, image_prompt, image_input],
                outputs=[image_result_image, image_result_text]
            )
            
            image_report_badcase_btn.click(
                fn=save_badcase_record,
                inputs=[image_prompt, image_input, image_result_text, image_result_text]
            ).then(
                fn=lambda: gr.Info("错误记录已保存"),
                inputs=[],
                outputs=[]
            )
        
        with gr.Tab("文字生成视频"):
            with gr.Row():
                with gr.Column():
                    video_prompt = gr.Textbox(label="提示词", lines=6, value="多个镜头。一名侦探进入一间光线昏暗的房间。他检查桌上的线索，手里拿起桌上的某个物品。镜头转向他正在思索。 --ratio 16:9")
                    video_generate_btn = gr.Button("生成视频")
                    video_report_badcase_btn = gr.Button("记录错误")
                    task_id = gr.Textbox(label="任务ID", visible=False)
                with gr.Column():
                    video_result_video = gr.Video(label="生成的视频")
                    video_result_text = gr.Textbox(label="结果详情", lines=15)
            
            video_generate_btn.click(
                fn=text_to_video,
                inputs=[api_key, video_prompt],
                outputs=[video_result_video, video_result_text, task_id]
            ).then(
                fn=auto_check_video_status,
                inputs=[api_key, task_id],
                outputs=[video_result_video, video_result_text, task_id]
            )
            
            video_report_badcase_btn.click(
                fn=save_badcase_record,
                inputs=[video_prompt, gr.State(None), video_result_text, video_result_text]
            ).then(
                fn=lambda: gr.Info("错误记录已保存"),
                inputs=[],
                outputs=[]
            )
        
        with gr.Tab("图片生成视频-首帧"):
            with gr.Row():
                with gr.Column():
                    image_to_video_ff_prompt = gr.Textbox(label="提示词", lines=6, value="保持商品的主体形态，商品为一个（填写自己的商品），摄像机缓慢地围绕产品进行360度轨道旋转，同时伴有轻微的推近和拉远运动，完美展示其前后左右所有角度和内部结构。镜头运动极其平滑，如丝般顺滑。--ratio adaptive  --dur 5")
                    image_to_video_ff_image = gr.Image(label="选择图片", type="pil")
                    image_to_video_ff_btn = gr.Button("生成视频")
                    image_to_video_ff_report_badcase_btn = gr.Button("记录错误")
                    image_to_video_ff_task_id = gr.Textbox(label="任务ID", visible=False)
                with gr.Column():
                    image_to_video_ff_result_video = gr.Video(label="生成的视频")
                    image_to_video_ff_result_text = gr.Textbox(label="结果详情", lines=15)
            
            image_to_video_ff_btn.click(
                fn=image_to_video_first_frame,
                inputs=[api_key, image_to_video_ff_prompt, image_to_video_ff_image],
                outputs=[image_to_video_ff_result_video, image_to_video_ff_result_text, image_to_video_ff_task_id]
            ).then(
                fn=auto_check_video_status,
                inputs=[api_key, image_to_video_ff_task_id],
                outputs=[image_to_video_ff_result_video, image_to_video_ff_result_text, image_to_video_ff_task_id]
            )
            
            image_to_video_ff_report_badcase_btn.click(
                fn=save_badcase_record,
                inputs=[image_to_video_ff_prompt, image_to_video_ff_image, image_to_video_ff_result_text, image_to_video_ff_result_text]
            ).then(
                fn=lambda: gr.Info("错误记录已保存"),
                inputs=[],
                outputs=[]
            )
        
        with gr.Tab("Badcase记录"):
            badcase_df = load_badcase_records()
            badcase_table = gr.DataFrame(label="Badcase记录", value=badcase_df)
            refresh_badcase_btn = gr.Button("刷新")
            
            refresh_badcase_btn.click(
                fn=load_badcase_records,
                inputs=[],
                outputs=badcase_table
            )
    # share=True
    demo.launch(pwa=True,share=True)

if __name__ == "__main__":
    main()